{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ISubpr_SSsiM"
      },
      "source": [
        "##### Copyright 2020 The TensorFlow Authors.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "3jTMb1dySr3V"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6DWfyNThSziV"
      },
      "source": [
        "# tf.function によるパフォーマンスの改善\n",
        "\n",
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>     <a target=\"_blank\" href=\"https://www.tensorflow.org/guide/function\" class=\"\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\">TensorFlow.org で表示</a>   </td>\n",
        "  <td>     <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs-l10n/blob/master/site/ja/guide/function.ipynb\" class=\"\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\">Google Colab で実行</a>   </td>\n",
        "  <td>     <a target=\"_blank\" href=\"https://github.com/tensorflow/docs-l10n/blob/master/site/ja/guide/function.ipynb\" class=\"\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\">GitHub でソースを表示</a>   </td>\n",
        "  <td>     <a href=\"https://storage.googleapis.com/tensorflow_docs/docs-l10n/site/ja/guide/function.ipynb\" class=\"\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\">ノートブックをダウンロード</a>   </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J122XQYG7W6w"
      },
      "source": [
        "TensorFlow 2 の Eager execution はデフォルトで有効になっています。ユーザーインターフェースは直感的で柔軟性に優れていますが（一度限りの演算の実行ははるかに簡単で高速に行われます）、パフォーマンスとデプロイ能力に影響がでることがあります。\n",
        "\n",
        "TensorFlow 2.0 では Eager Execution の使いやすさとTensorFlow 1.0 のパワーとを同時に提供します。この統合の中核となるのは `tf.function` です。これは Python の構文のサブセットを移植可能でハイパフォーマンスな TensorFlow のグラフに変換します。\n",
        "\n",
        "このチュートリアルでは `tf.function` と AutoGraph の基本的な特徴についてひととおり確認します。\n",
        "\n",
        "主に次の内容と推奨事項について説明しています。\n",
        "\n",
        "- Eager モードでデバッグしてから、`@tf.function` でデコレートする。\n",
        "- オブジェクトミューテーションまたはリストの追加といった Python 側の効果に依存しないこと。\n",
        "- `tf.function` は TensorFlow 演算子と最も相性が良く、NumPy と Python 呼び出しは定数に変換される。\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SjvqpgepHJPd"
      },
      "source": [
        "## セットアップ"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "otIdN1TS8N7S"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I0xDjO4SHLUD"
      },
      "source": [
        "発生する可能性のあるエラーの種類を示すヘルパー関数を定義します。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "D25apou9IOXa"
      },
      "outputs": [],
      "source": [
        "import traceback\n",
        "import contextlib\n",
        "\n",
        "# Some helper code to demonstrate the kinds of errors you might encounter.\n",
        "@contextlib.contextmanager\n",
        "def assert_raises(error_class):\n",
        "  try:\n",
        "    yield\n",
        "  except error_class as e:\n",
        "    print('Caught expected exception \\n  {}:'.format(error_class))\n",
        "    traceback.print_exc(limit=2)\n",
        "  except Exception as e:\n",
        "    raise e\n",
        "  else:\n",
        "    raise Exception('Expected {} to be raised but no error was raised!'.format(\n",
        "        error_class))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WPSfepzTHThq"
      },
      "source": [
        "## 基礎"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CNwYTIJ8r56W"
      },
      "source": [
        "### 使い方\n",
        "\n",
        "定義する `Function`（`@tf.function` デコレーターを適用するなどして）は、コアの TensorFlow 演算とまったく変わりません。Eager での実行や勾配の計算などを行えます。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SbtT1-Wm70F2"
      },
      "outputs": [],
      "source": [
        "@tf.function  # The decorator converts `add` into a `Function`.\n",
        "def add(a, b):\n",
        "  return a + b\n",
        "\n",
        "add(tf.ones([2, 2]), tf.ones([2, 2]))  #  [[2., 2.], [2., 2.]]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uP-zUelB8DbX"
      },
      "outputs": [],
      "source": [
        "v = tf.Variable(1.0)\n",
        "with tf.GradientTape() as tape:\n",
        "  result = add(v, 1.0)\n",
        "tape.gradient(result, v)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ocWZvqrmHnmX"
      },
      "source": [
        "`Function` をほかの `Function` 内で使用できます。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "l5qRjdbBVdU6"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def dense_layer(x, w, b):\n",
        "  return add(tf.matmul(x, w), b)\n",
        "\n",
        "dense_layer(tf.ones([3, 2]), tf.ones([2, 2]), tf.ones([2]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "piBhz7gYsHqU"
      },
      "source": [
        "`Function` は、特に小さな演算が多数含まれるグラフでは、Eager コードよりも高速に実行されることがありますが、高価な演算がいくつか含まれるグラフ（畳み込みなど）では、速度の差はあまり見られません。\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zuXt4wRysI03"
      },
      "outputs": [],
      "source": [
        "import timeit\n",
        "conv_layer = tf.keras.layers.Conv2D(100, 3)\n",
        "\n",
        "@tf.function\n",
        "def conv_fn(image):\n",
        "  return conv_layer(image)\n",
        "\n",
        "image = tf.zeros([1, 200, 200, 100])\n",
        "# warm up\n",
        "conv_layer(image); conv_fn(image)\n",
        "print(\"Eager conv:\", timeit.timeit(lambda: conv_layer(image), number=10))\n",
        "print(\"Function conv:\", timeit.timeit(lambda: conv_fn(image), number=10))\n",
        "print(\"Note how there's not much difference in performance for convolutions\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uZ4Do2AV80cO"
      },
      "source": [
        "### トレーシング\n",
        "\n",
        "このセクションは、`Function` の内部動作や実装の詳細を説明します。*将来的に変更する可能性があります*が、いつなぜトレーシングが発生するのかを理解しておけば、`tf.function` を効果的に使用しやすくなります。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nhpUtRqsXoyM"
      },
      "source": [
        "#### 「トレーシング」とは？\n",
        "\n",
        "`Function` は [TensorFlow Graph](https://www.tensorflow.org/guide/intro_to_graphs#what_are_graphs) でプログラムを実行しますが、`tf.Graph` は、Eager TensorFlow プログラムにユーザーが記述するすべてのものを表現することはできません。たとえば、Python はポリモーフィズムをサポートしていますが、`tf.Graph` では、その入力に特定のデータ型と次元が必要です。またはコマンドラインの引数を読み取る、エラーを発生させる、より複雑な Python オブジェクトを扱うといったサイドタスクを実施しようとしても、どれも `tf.Graph` で実行することはできません。\n",
        "\n",
        "`Function` はコードを 2 つの段階に分けることで、このギャップの橋渡しの役割を果たします。\n",
        "\n",
        "1. 「**トレーシング**」と呼ばれる第 1 段階において、`Function` は新しい `tf.Graph` を作成します。Python コードは通常通り実行しますが、すべての TensorFlow 演算（2 つのテンソルを加算するなど）は *据え置き*となります。これらは `tf.Graph` にとらわれるため、実行しません。\n",
        "\n",
        "2. 第 2 段階では、最初の段階で据え置きとなったすべての演算を含む `tf.Graph` が実行されます。この段階は、トレーシングの段階よりもはるかに高速に行われます。\n",
        "\n",
        "`Function` は、その入力によっては必ずしも最初の段階で呼び出されたときに実行するわけではありません。この判定がどのように行われるのかについては、以下の「[トレーシングの規則](#rules_of_tracing)」をご覧ください。最初の段階を省略して 2 番目の段階のみを実行できれば、TensorFlow の高いパフォーマンスが発揮されます。\n",
        "\n",
        "`Function` がトレーシングしないと判断した場合、トレーシング段階の直後に 第 2 段階が始まるため、`Function` を呼び出すと、`tf.Graph` の作成と実行が行われます。後の方で、[`get_concrete_function`](#obtaining_concrete_functions) を使ってトレーシング段階のみを実行する方法を説明します。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "K7scSzLx662f"
      },
      "source": [
        "型の異なる引数を `Function` に渡すと、両方の段階が実行されます。\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kojmJrgq8U9v"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def double(a):\n",
        "  print(\"Tracing with\", a)\n",
        "  return a + a\n",
        "\n",
        "print(double(tf.constant(1)))\n",
        "print()\n",
        "print(double(tf.constant(1.1)))\n",
        "print()\n",
        "print(double(tf.constant(\"a\")))\n",
        "print()\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QPfouGUQrcNb"
      },
      "source": [
        "同じ型の引数で `Function` を繰り返し呼び出すと、生成されるグラフはまったく同じになるため、TensorFlow はトレーシング段階を省略して前にトレーシングしたグラフを再利用することに注意してください。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hFccbWFRrsBp"
      },
      "outputs": [],
      "source": [
        "# This doesn't print 'Tracing with ...'\n",
        "print(double(tf.constant(\"b\")))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fgIO_XEzcB9o"
      },
      "source": [
        "すべての利用可能なトレースを確認するには、`pretty_printed_concrete_signatures()` を使用できます。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IiQc4IKAb-NX"
      },
      "outputs": [],
      "source": [
        "print(double.pretty_printed_concrete_signatures())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rKQ92VEWI7n8"
      },
      "source": [
        "ここまで、`tf.function` が TensorFlow のグラフトレーシングロジックにキャッシュされた動的ディスパッチレイヤーを作成するのを見てきました。用語についてより具体的に説明すると、次のように言えます。\n",
        "\n",
        "- `tf.Graph` は、言語に依存しない、生の移植可能な TensorFlow 計算の表現です。\n",
        "- `ConcreteFunction` は `tf.Graph` をラップします。\n",
        "- `Function` は `ConcreteFunction` のキャッシュを管理し、入力に適したものを選択します。\n",
        "- `tf.function` は Python 関数をラップし、`Function` オブジェクトを返します。\n",
        "- **トレーシング**は `tf.Graph` を作成し、それを `ConcreteFunction`（または**トレース**）をラップします。\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "129-iRsPS-gY"
      },
      "source": [
        "#### トレーシングの規則\n",
        "\n",
        "`Function` は、トレーシングされた `ConcreteFunction` を再利用するかどうかを判定します。この判定は、入力の引数とキーワードから**キャッシュキー**を計算して行われます。**キャッシュキー**は、次の規則（変更される可能性があります）に従って、`Function` 呼び出しの入力引数とキーワードに基づく `ConcreteFunction` を識別するキーです。\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h62XoXho6EWN"
      },
      "source": [
        "- `tf.Tensor` に生成されたキーは、その形状と dtype である。\n",
        "- `tf.Variable` に生成されたキーは、一意の変数 ID である。\n",
        "- Python プリミティブ型（`int`、`float`、`str` など）に生成されたキーはその値である。\n",
        "- ネストされた `dict`、`list`、`tuple`、`namedtuple`、および [`attr`](https://www.attrs.org/en/stable/) に生成されたキーは、フラット化されたリーフキーのタプルである（`nest.flatten` を参照）。（このフラット化の結果、ネスト構造がトレーシング時とは異なる具象関数が呼び出されると、TypeError が発生します。）\n",
        "- そのたすべての Python の型については、キーはオブジェクト固有です。このため、関数またはメソッドは呼び出しに使用されるインスタンスごとにトレースされます。\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GNNN4lgRzpIs"
      },
      "source": [
        "注意: キャッシュキーは、`Function` の入力パラメータに基づくため、グローバルと[自由変数](https://docs.python.org/3/reference/executionmodel.html#binding-of-names)を変更するだけでは、新しいトレースは作成されません。Python のグローバル変数と自由変数を扱う際の推奨される方法については、[こちらのセクション](#depending_on_python_global_and_free_variables)をご覧ください。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PEDwbumO32Wh"
      },
      "source": [
        "#### リトレーシングの制御\n",
        "\n",
        "リトレーシングは、`Function` が 2 つ以上のトレースを作成する際に発生します。これは、TensorFlow が一連の入力ごとに正しいグラフを生成する上で役立ちますが、トレーシングは高価な演算です。`Function` が呼び出しごとに新しいグラフをリトレーシングすると、コードの実行は `tf.function` を使用しない場合よりも遅くなってしまいます。\n",
        "\n",
        "トレーシングの動作を制御するには、次のテクニックを使用できます。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EUtycWJa34TT"
      },
      "source": [
        "- トレーシングを制限するために、`input_signature` を `tf.function` に指定します。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_BDMIRmu1RGB"
      },
      "outputs": [],
      "source": [
        "@tf.function(input_signature=(tf.TensorSpec(shape=[None], dtype=tf.int32),))\n",
        "def next_collatz(x):\n",
        "  print(\"Tracing with\", x)\n",
        "  return tf.where(x % 2 == 0, x // 2, 3 * x + 1)\n",
        "\n",
        "print(next_collatz(tf.constant([1, 2])))\n",
        "# We specified a 1-D tensor in the input signature, so this should fail.\n",
        "with assert_raises(ValueError):\n",
        "  next_collatz(tf.constant([[1, 2], [3, 4]]))\n",
        "\n",
        "# We specified an int32 dtype in the input signature, so this should fail.\n",
        "with assert_raises(ValueError):\n",
        "  next_collatz(tf.constant([1.0, 2.0]))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ocxX-HVk7P2o"
      },
      "source": [
        "- トレースを柔軟に再利用できるようにするために、[None] 次元を `tf.TensorSpec` に指定します。\n",
        "\n",
        "    TensorFlow は形状に基づいてテンソルを一致させるため、ワイルドカードとして `None` 次元を使用することで、`Function` が可変サイズの入力にトレースを再利用できるようになります。可変サイズの入力は、長さの異なるシーケンスがある場合や、バッチごとに画像のサイズが異なる場合に発生します（例として、[Transformer](../tutorials/text/transformer.ipynb) と [Deep Dream](../tutorials/generative/deepdream.ipynb) チュートリアルをご覧ください）。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4Viun7dh7PmF"
      },
      "outputs": [],
      "source": [
        "@tf.function(input_signature=(tf.TensorSpec(shape=[None], dtype=tf.int32),))\n",
        "def g(x):\n",
        "  print('Tracing with', x)\n",
        "  return x\n",
        "\n",
        "# No retrace!\n",
        "print(g(tf.constant([1, 2, 3])))\n",
        "print(g(tf.constant([1, 2, 3, 4, 5])))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AY5oiQN0XIyA"
      },
      "source": [
        "- リトレーシングを減らすために、テンソルに Python 引数をキャストします。\n",
        "\n",
        "    通常、Python 引数は、`num_layers=10` または `training=True` または `nonlinearity='relu'` などのように、ハイパーパラメータとグラフ構造の制御に使用されます。そのため、Python 引数が変わると、当然グラフをリトレースする必要が出てきます。\n",
        "\n",
        "    しかし、Python 引数がグラフ構造の制御に使用されていない場合もあります。こういった場合、Python の値の変化によってリトレーシングがトリガーされますが、これは不要です。この、AutoGraph が動的にアンロールするトレーニングループを例に見てみましょう。トレースが何度も行われますが、生成されたグラフはまったく同じであるため、リトレーシングは不要と言えます。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uydzR5JYUU8H"
      },
      "outputs": [],
      "source": [
        "def train_one_step():\n",
        "  pass\n",
        "\n",
        "@tf.function\n",
        "def train(num_steps):\n",
        "  print(\"Tracing with num_steps = \", num_steps)\n",
        "  tf.print(\"Executing with num_steps = \", num_steps)\n",
        "  for _ in tf.range(num_steps):\n",
        "    train_one_step()\n",
        "\n",
        "print(\"Retracing occurs for different Python arguments.\")\n",
        "train(num_steps=10)\n",
        "train(num_steps=20)\n",
        "\n",
        "print()\n",
        "print(\"Traces are reused for Tensor arguments.\")\n",
        "train(num_steps=tf.constant(10))\n",
        "train(num_steps=tf.constant(20))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4pJqkDR_Q2wz"
      },
      "source": [
        "リトレーシングを強制する必要がある場合は、新しい `Function` を作成します。トレースは絶対に、各 `Function` オブジェクト間で共有されることはありません。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uHp4ousu4DdN"
      },
      "outputs": [],
      "source": [
        "def f():\n",
        "  print('Tracing!')\n",
        "  tf.print('Executing')\n",
        "\n",
        "tf.function(f)()\n",
        "tf.function(f)()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "96IxS2WR37fF"
      },
      "source": [
        "### 具象関数の取得\n",
        "\n",
        "関数がトレースされるたびに新しい具象関数が作成されますが、`get_concrete_function` を使うことで、具象関数を直接取得できます。\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mHg2CGtPQ3Hz"
      },
      "outputs": [],
      "source": [
        "print(\"Obtaining concrete trace\")\n",
        "double_strings = double.get_concrete_function(tf.constant(\"a\"))\n",
        "print(\"Executing traced function\")\n",
        "print(double_strings(tf.constant(\"a\")))\n",
        "print(double_strings(a=tf.constant(\"b\")))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6IVZ-NVf9vsx"
      },
      "outputs": [],
      "source": [
        "# You can also call get_concrete_function on an InputSpec\n",
        "double_strings_from_inputspec = double.get_concrete_function(tf.TensorSpec(shape=[], dtype=tf.string))\n",
        "print(double_strings_from_inputspec(tf.constant(\"c\")))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iR4fVmG34xvF"
      },
      "source": [
        "`ConcreteFunction` を出力すると、入力引数（型付き）とその出力型の概要が表示されます。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "o3-JbkIk41r8"
      },
      "outputs": [],
      "source": [
        "print(double_strings)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QtqfvljZeuOV"
      },
      "source": [
        "また、具象関数のシグネチャを直接取得することもできます。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nzbrqFABe0zG"
      },
      "outputs": [],
      "source": [
        "print(double_strings.structured_input_signature)\n",
        "print(double_strings.structured_outputs)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lar5A_5m5IG1"
      },
      "source": [
        "互換性のない型で具象トレースを使用すると、エラーが発生します。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "G5eeTK-T5KYj"
      },
      "outputs": [],
      "source": [
        "with assert_raises(tf.errors.InvalidArgumentError):\n",
        "  double_strings(tf.constant(1))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "st2L9VNQVtSG"
      },
      "source": [
        "Python 引数は、具象関数の入力シグネチャで特別に扱われていることに気づいたかもしれません。TensorFlow 2.3 より前では、Python 引数は単に具象関数のシグネチャから削除されていましたが、TensorFlow 2.3 からはシグネチャに残されたまま、トレーシング中に値セットを取るように制約されています。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "U_QyPSGoaC35"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def pow(a, b):\n",
        "  return a ** b\n",
        "\n",
        "square = pow.get_concrete_function(a=tf.TensorSpec(None, tf.float32), b=2)\n",
        "print(square)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "E76vIDhQbXIb"
      },
      "outputs": [],
      "source": [
        "assert square(tf.constant(10.0)) == 100\n",
        "\n",
        "with assert_raises(TypeError):\n",
        "  square(tf.constant(10.0), b=3)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "41gJh_JGIfuA"
      },
      "source": [
        "### グラフの取得\n",
        "\n",
        "それぞれの具象関数は、`tf.Graph` を囲む呼び出し可能なラッパーです。通常、実際の `tf.Graph` オブジェクトを取得する必要はないにしろ、具象関数から簡単に取得することが可能です。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5UENeGHfaX8g"
      },
      "outputs": [],
      "source": [
        "graph = double_strings.graph\n",
        "for node in graph.as_graph_def().node:\n",
        "  print(f'{node.input} -> {node.name}')\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aIKkgr6qdtp4"
      },
      "source": [
        "### デバッグ\n",
        "\n",
        "一般的に、コードのデバックは、`tf.function` 内で行うよりも、Eager モードで行う方が簡単です。Eager モードでは、`tf.function` でデコレートする前に、コードがエラーなく実行することを確認しておく必要があります。デバッグプロセスを支援する目的で、`tf.config.run_functions_eagerly(True)` を呼び出すと、`tf.function` をグローバルに無効にして、有効にし直すことができます。\n",
        "\n",
        "`tf.function` 内でのみ出現する問題を追跡する場合、次のようなヒントがあります。\n",
        "\n",
        "- 従来のシンプルな Python `print` 呼び出しは、トレーシング中にのみ実行されるため、関数が（リ）トレーシングされるときに追跡しやすくなります。\n",
        "- `tf.print` 呼び出しは毎回実行するため、実行中の中間値の追跡に役立ちます。\n",
        "- `tf.debugging.enable_check_numerics` は、NaN と Inf がいつ作成されるかを簡単に追跡できます。\n",
        "- `pdb` は、トレーシング中に何が起きているのかを理解する上で役立ちます。（注意: PDB が示すのは、AutoGraph 変換ソースコードです。）"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5f05Vr_YBUCz"
      },
      "source": [
        "## AutoGraph 変換\n",
        "\n",
        "AutoGraph は、`tf.function` 内でデフォルトで利用できるようになっているライブラリで、Python の Eager コードのサブセットとグラフ対応の TensorFlow 演算に変換します。これには、`if`、`for`、`while` などの制御フローが含まれます。\n",
        "\n",
        "`tf.cond` や `tf.while_loop` などの TensorFlow 演算は機能し続けますが、制御フローは、Python で記述された場合の方が書きやすく理解しやすいことがほとんどです。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yCQTtTPTW3WF"
      },
      "outputs": [],
      "source": [
        "# Simple loop\n",
        "\n",
        "@tf.function\n",
        "def f(x):\n",
        "  while tf.reduce_sum(x) > 1:\n",
        "    tf.print(x)\n",
        "    x = tf.tanh(x)\n",
        "  return x\n",
        "\n",
        "f(tf.random.uniform([5]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KxwJ8znPI0Cg"
      },
      "source": [
        "興味があれば、AutoGraph が生成するコードを検査できます。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jlQD1ffRXJhl"
      },
      "outputs": [],
      "source": [
        "print(tf.autograph.to_code(f.python_function))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xgKmkrNTZSyz"
      },
      "source": [
        "### 条件文\n",
        "\n",
        "AutoGraph は `if <condition>` 文を相当する `tf.cond` 呼び出しに変換します。この置換は、`<condition>` がテンソルである場合に行われます。テンソルでない場合は、`if` 文は Python の条件文として実行されます。\n",
        "\n",
        "Python 条件文はトレーシング中に実行するため、条件文のブランチが 1 つだけグラフに追加されます。AutoGraph を使用しない場合、データに依存する制御フローが存在すると、トレーシングされたこのグラフは別のブランチを取ることができません。\n",
        "\n",
        "`tf.cond` は、条件文の両方のブランチをトレーシングし、実行時に動的に 1 つのブランチを選択してグラフに追加します。トレーシングには意図しない副作用がある場合があります。詳細は、[AutoGraph のトレーシング効果](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/autograph/g3doc/reference/control_flow.md#effects-of-the-tracing-process)をご覧ください。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BOQl8PMq2Sf3"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def fizzbuzz(n):\n",
        "  for i in tf.range(1, n + 1):\n",
        "    print('Tracing for loop')\n",
        "    if i % 15 == 0:\n",
        "      print('Tracing fizzbuzz branch')\n",
        "      tf.print('fizzbuzz')\n",
        "    elif i % 3 == 0:\n",
        "      print('Tracing fizz branch')\n",
        "      tf.print('fizz')\n",
        "    elif i % 5 == 0:\n",
        "      print('Tracing buzz branch')\n",
        "      tf.print('buzz')\n",
        "    else:\n",
        "      print('Tracing default branch')\n",
        "      tf.print(i)\n",
        "\n",
        "fizzbuzz(tf.constant(5))\n",
        "fizzbuzz(tf.constant(20))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4rBO5AQ15HVC"
      },
      "source": [
        "AutoGraph 変換の if 文におけるその他の制約事項については、[リファレンスドキュメント](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/autograph/g3doc/reference/control_flow.md#if-statements)をご覧ください。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yho4J0a0ZkQS"
      },
      "source": [
        "### ループ\n",
        "\n",
        "AutoGraph は、一部の `for` 文と `while` 文を相当する `tf.while_loop` などの TensorFlow のループ演算に変換します。変換されない場合、`for` または `while` ループは Python ループとして実行されます。\n",
        "\n",
        "この置き換えは、次の場合に行われます。\n",
        "\n",
        "- `for x in y`: `y` がテンソルである場合、`tf.while_loop` に変換されます。`y` が `tf.data.Dataset` である特別なケースでは、`tf.data.Dataset` 演算の組み合わせが生成されます。\n",
        "- `while <condition>`: `<condition>` がテンソルである場合、`tf.while_loop` に変換されます。\n",
        "\n",
        "Python ループは、トレーシング中に実行され、ループの反復ごとに、`tf.Graph` に追加の演算が追加されます。\n",
        "\n",
        "TensorFlow ループはループの本体をトレーシングし、実行時に実行する反復回数を動的に選択します。ループ本体は、生成された `tf.Graph` に一度だけ出現します。\n",
        "\n",
        "AutoGraph 変換の `for` 文と `while` 文におけるその他の制約事項については、[リファレンスドキュメント](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/autograph/g3doc/reference/control_flow.md#while-statements)をご覧ください。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sp4rbIdfbM6s"
      },
      "source": [
        "#### Python データのループ\n",
        "\n",
        "一般的な落とし穴は、`tf.function` 内で Python/Numpy データをループする際にあります。このループは、トレーシングプロセス中に実行し、ループの反復ごとに<br>モデルのコピーを `tf.Graph` に追加してしまいます。\n",
        "\n",
        "トレーニングループ全体を `tf.function` にラップしたいのであれば、データを `tf.data.Dataset` としてラップし、AutoGraph にトレーニングループを動的に展開させるようにするのが最も安全な方法です。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WGZ19LspbZ27"
      },
      "outputs": [],
      "source": [
        "def measure_graph_size(f, *args):\n",
        "  g = f.get_concrete_function(*args).graph\n",
        "  print(\"{}({}) contains {} nodes in its graph\".format(\n",
        "      f.__name__, ', '.join(map(str, args)), len(g.as_graph_def().node)))\n",
        "\n",
        "@tf.function\n",
        "def train(dataset):\n",
        "  loss = tf.constant(0)\n",
        "  for x, y in dataset:\n",
        "    loss += tf.abs(y - x) # Some dummy computation.\n",
        "  return loss\n",
        "\n",
        "small_data = [(1, 1)] * 3\n",
        "big_data = [(1, 1)] * 10\n",
        "measure_graph_size(train, small_data)\n",
        "measure_graph_size(train, big_data)\n",
        "\n",
        "measure_graph_size(train, tf.data.Dataset.from_generator(\n",
        "    lambda: small_data, (tf.int32, tf.int32)))\n",
        "measure_graph_size(train, tf.data.Dataset.from_generator(\n",
        "    lambda: big_data, (tf.int32, tf.int32)))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JeD2U-yrbfVb"
      },
      "source": [
        "Python/Numpy データを Dataset にラップする際は、`tf.data.Dataset.from_generator` と ` tf.data.Dataset.from_tensors` の違いに注意してください。前者は、データを Python に維持し、`tf.py_function` 経由で取得するため、パフォーマンスに問題がありますが、後者は、データのコピーをグラフ内の大型の `tf.constant()` ノードとしてバンドル化するため、メモリに問題が現れます。\n",
        "\n",
        "データを消費するには、TFRecordDataset/CsvDataset などでファイルからデータを読み取るのが最も効果的な方法です。そうすれば、Python を使わずに、TensorFlow 自体でデータの非同期読み込みとプリフェッチを管理できるようになります。詳細は、[tf.data guide](../../guide/data) をご覧ください。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hyksHW9TCukR"
      },
      "source": [
        "#### ループでの値の累積\n",
        "\n",
        "ループの反復ごとに値を累積していくのは一般的なパターンです。通常は、Python のリストに追加したり、Python ディレクトリにエントリを追加したりして行われますが、これらは Python の副作用であるため、動的に展開されるループでは期待どおりに動作しません。動的に展開されるループの結果を累積する場合は、`tf.TensorArray` を使用してください。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HJ3Vb3dXfefN"
      },
      "outputs": [],
      "source": [
        "batch_size = 2\n",
        "seq_len = 3\n",
        "feature_size = 4\n",
        "\n",
        "def rnn_step(inp, state):\n",
        "  return inp + state\n",
        "\n",
        "@tf.function\n",
        "def dynamic_rnn(rnn_step, input_data, initial_state):\n",
        "  # [batch, time, features] -> [time, batch, features]\n",
        "  input_data = tf.transpose(input_data, [1, 0, 2])\n",
        "  max_seq_len = input_data.shape[0]\n",
        "\n",
        "  states = tf.TensorArray(tf.float32, size=max_seq_len)\n",
        "  state = initial_state\n",
        "  for i in tf.range(max_seq_len):\n",
        "    state = rnn_step(input_data[i], state)\n",
        "    states = states.write(i, state)\n",
        "  return tf.transpose(states.stack(), [1, 0, 2])\n",
        "  \n",
        "dynamic_rnn(rnn_step,\n",
        "            tf.random.uniform([batch_size, seq_len, feature_size]),\n",
        "            tf.zeros([batch_size, feature_size]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i2MVoIVaNApG"
      },
      "source": [
        "## 制限事項\n",
        "\n",
        "TensorFlow の `Function` には、設計上、いくつかの制限事項があり、Python 関数を `Function` に変換する際には、注意が必要です。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EJqHGFSVLIKl"
      },
      "source": [
        "### Python の副作用の実行\n",
        "\n",
        "`Function` 内での出力、リストへのアペンド、グローバル変数のミューテーションといった副作用は、2 回実行されたり、まったく実行しなかったりといったように、予測のつかない動作をすることがあります。また、入力セットで `Function` を初めて呼び出した場合にのみ実行し、以降では、Python コードを実行せずに、トレーシング済みの `tf.Graph` が再実行されてしまうこともあります。\n",
        "\n",
        "基本的に、ロジックでは Python の副作用に依存しないようにし、トレースをデバッグするためだけに使用することをお勧めします。呼び出しごとに TensorFlow ランタイムが確実にコードを実行できるようにするには、`tf.data`、`tf.print`、`tf.summary`、`tf.Variable.assign`、`tf.TensorArray` などの TensorFlow API を使用するのが最善の方法です。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "w2sACuZ9TTRk"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def f(x):\n",
        "  print(\"Traced with\", x)\n",
        "  tf.print(\"Executed with\", x)\n",
        "\n",
        "f(1)\n",
        "f(1)\n",
        "f(2)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e1I0dPiqTV8H"
      },
      "source": [
        "`Function` の呼び出しごとに Python コードを実行する場合は、`tf.py_function` が脱出口です。code2}tf.py_function には移植性がなく、特にパフォーマンスに優れているわけでもなく、SavedModel で保存できなければ、分散型（マルチ GPU、TPU）の環境でうまく動作するわけでもありません。また、`tf.py_function` はグラフに組み込む必要もあるため、すべての入力/出力をテンソルにキャストしてしまいます。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bOW1v9WVKGgH"
      },
      "source": [
        "#### Python のグローバル変数と自由変数の変更\n",
        "\n",
        "Python のグローバル変数と[自由変数](https://docs.python.org/3/reference/executionmodel.html#binding-of-names)の変更は、Python の副作用としてみなされるため、トレーシング中にのみ発生します。\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7aJD--9qTWmg"
      },
      "outputs": [],
      "source": [
        "external_list = []\n",
        "\n",
        "@tf.function\n",
        "def side_effect(x):\n",
        "  print('Python side effect')\n",
        "  external_list.append(x)\n",
        "\n",
        "side_effect(1)\n",
        "side_effect(1)\n",
        "side_effect(1)\n",
        "# The list append only happened once!\n",
        "assert len(external_list) == 1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pbFG5CX4LwQA"
      },
      "source": [
        "リスト、辞書、`Function` の外側で機能するその他のオブジェクトなどのコンテナのミューテーションは避けてください。代わりに、引数と TF オブジェクトを使用しましょう。たとえば、「[ループでの値の累積](#accumulating_values_in_a_loop)」セクションには、リストのような演算を実装する方法の一例が示されています。\n",
        "\n",
        "一部のケースでは、[`tf.Variable`](https://www.tensorflow.org/guide/variable) である場合に状態をキャプチャして操作することができます。Keras モデルの重みは、このようにして、同じ `ConcreteFunction` への呼び出しの繰り返しで更新されています。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "X_oNNGrAqPJ1"
      },
      "source": [
        "#### Python イテレータとジェネレータの使用"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "msTmv-oyUNaf"
      },
      "source": [
        "ジェネレータやイテレータなどの多くの Python 機能は、Python ランタイムに依存して状態を追跡しています。一般的に、これらのコンストラクトは Eager モードでも期待どおりに動作しますが、Python の副作用の例であるため、トレーシング中にしか発生しません。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FNPD4unZUedH"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def buggy_consume_next(iterator):\n",
        "  tf.print(\"Value:\", next(iterator))\n",
        "\n",
        "iterator = iter([1, 2, 3])\n",
        "buggy_consume_next(iterator)\n",
        "# This reuses the first value from the iterator, rather than consuming the next value.\n",
        "buggy_consume_next(iterator)\n",
        "buggy_consume_next(iterator)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wcS3TAgCjTWR"
      },
      "source": [
        "TensorFlow にリストコントラクト用の特別な `tf.TensorArray` があるように、イテレーション用にも特別な `tf.data.Iterator` があります。概要は、[AutoGraph 変換](#autograph_transformations)をご覧ください。また、[`tf.data`](https://www.tensorflow.org/guide/data) API を使って、ジェネレータのパターンを実装できます。\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8D_iKetXW6VE"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def good_consume_next(iterator):\n",
        "  # This is ok, iterator is a tf.data.Iterator\n",
        "  tf.print(\"Value:\", next(iterator))\n",
        "\n",
        "ds = tf.data.Dataset.from_tensor_slices([1, 2, 3])\n",
        "iterator = iter(ds)\n",
        "good_consume_next(iterator)\n",
        "good_consume_next(iterator)\n",
        "good_consume_next(iterator)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FHQ0UeU-vWo8"
      },
      "source": [
        "### `Function` 呼び出し間での tf.Variables の削除\n",
        "\n",
        "発生する可能性のあるもう 1 つのエラーは、ガベージコレクションされた変数です。`ConcreteFunction` は使用した変数への [WeakRefs](https://docs.python.org/3/library/weakref.html) のみを保持するため、ユーザーが変数への参照を保持するようにする必要があります。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uMiRPfETjpt-"
      },
      "outputs": [],
      "source": [
        "external_var = tf.Variable(3)\n",
        "@tf.function\n",
        "def f(x):\n",
        "  return x * external_var\n",
        "\n",
        "traced_f = f.get_concrete_function(4)\n",
        "print(\"Calling concrete function...\")\n",
        "print(traced_f(4))\n",
        "\n",
        "# The original variable object gets garbage collected, since there are no more\n",
        "# references to it.\n",
        "external_var = tf.Variable(4)\n",
        "print()\n",
        "print(\"Calling concrete function after garbage collecting its closed Variable...\")\n",
        "with assert_raises(tf.errors.FailedPreconditionError):\n",
        "  traced_f(4)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-D6nh3QirXAd"
      },
      "source": [
        "## 既知の問題\n",
        "\n",
        "`Function` が正しく評価していない場合、以下の既知の問題が該当する可能性があります。これらの問題は、今後修正される予定です。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZoPg5w1Pjqna"
      },
      "source": [
        "### Python のグローバル変数と自由変数への依存\n",
        "\n",
        "`Function` は、Python 引数の新しい値で呼び出された時に新しい `ConcreteFunction` を作成しますが、Python クロージャ、グローバル変数、またはその `Function` の非ローカル変数に対しては作成しません。`Function` への呼び出しごとに値が変化する場合でも、`Function` はトレーシングされたときの値をそのまま使用してしまいます。これは、通常の Python 関数の動作とは異なります。\n",
        "\n",
        "このため、外側の名前を閉じる代わりに引数を使用する関数プログラミングの様式をお勧めします。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oeJMdXd3M0cM"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def buggy_add():\n",
        "  return 1 + foo\n",
        "\n",
        "@tf.function\n",
        "def recommended_add(foo):\n",
        "  return 1 + foo\n",
        "\n",
        "foo = 1\n",
        "print(\"Buggy:\", buggy_add())\n",
        "print(\"Correct:\", recommended_add(foo))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "L3q7sUJWZOSU"
      },
      "outputs": [],
      "source": [
        "print(\"Updating the value of `foo` to 100!\")\n",
        "foo = 100\n",
        "print(\"Buggy:\", buggy_add())  # Did not change!\n",
        "print(\"Correct:\", recommended_add(foo))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tu0SnPwaL7pI"
      },
      "source": [
        "外側の名前は、その値を更新しない場合にのみ閉じることができます。\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hvwe9gTIWfx6"
      },
      "source": [
        "#### Python オブジェクトへの依存"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BJkZS-SwPvOQ"
      },
      "source": [
        "Python オブジェクトを引数として `tf.function` に渡す上での推奨事項には、多数の既知の問題があります。これらは今後修正される予定です。一般的に、Python のプリミティブ型または `tf.nest` と互換性のある構造を引数として使用する場合や、オブジェクトの*別の*インスタンスを `Function` に渡す場合には、一貫したトレーシングを期待できますが、**同一のオブジェクトであっても属性が異なる**ものを渡す場合、`Function` は、新しいトレースを*作成しません*。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ux8KJESVWDxX"
      },
      "outputs": [],
      "source": [
        "class SimpleModel(tf.Module):\n",
        "  def __init__(self):\n",
        "    # These values are *not* tf.Variables.\n",
        "    self.bias = 0.\n",
        "    self.weight = 2.\n",
        "\n",
        "@tf.function\n",
        "def evaluate(model, x):\n",
        "  return model.weight * x + model.bias\n",
        "\n",
        "simple_model = SimpleModel()\n",
        "x = tf.constant(10.)\n",
        "print(evaluate(simple_model, x))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mUxRF4ghZZvX"
      },
      "outputs": [],
      "source": [
        "print(\"Adding bias!\")\n",
        "simple_model.bias += 5.0\n",
        "print(evaluate(simple_model, x))  # Didn't change :("
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ytcgg2qFWaBF"
      },
      "source": [
        "同じ `Function` を使用して、更新されたモデルのインスタンスを評価する場合、更新されたモデルには、元のモデルと[同じキャッシュキー](#rules_of_tracing)が含まれるため、不具合が生じます。\n",
        "\n",
        "このため、ミュート可能なオブジェクト属性に依存しない `Function` を記述するか、新しいオブジェクトを作成することをお勧めします。\n",
        "\n",
        "この方法が困難な場合は、回避策として、オブジェクトを変更するたびに新しい `Function` がリトレーシングを行うようにする方法が挙げられます。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pFvWmWAAQjrv"
      },
      "outputs": [],
      "source": [
        "def evaluate(model, x):\n",
        "  return model.weight * x + model.bias\n",
        "\n",
        "new_model = SimpleModel()\n",
        "evaluate_no_bias = tf.function(evaluate).get_concrete_function(new_model, x)\n",
        "# Don't pass in `new_model`, `Function` already captured its state during tracing.\n",
        "print(evaluate_no_bias(x))  "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bdU2-jF4ZH0B"
      },
      "outputs": [],
      "source": [
        "print(\"Adding bias!\")\n",
        "new_model.bias += 5.0\n",
        "# Create new Function and ConcreteFunction since you modified new_model.\n",
        "evaluate_with_bias = tf.function(evaluate).get_concrete_function(new_model, x)\n",
        "print(evaluate_with_bias(x)) # Don't pass in `new_model`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uFgEZClsZrEi"
      },
      "source": [
        "[リトレーシングにはコストがかかる](https://www.tensorflow.org/guide/intro_to_graphs#tracing_and_performance)ため、`tf.Variable` をオブジェクト属性として使用することができます。こうすることで、リトレーシングを行わずに、ミュートして（変更はしません）同様の効果を得ることができます。\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "daAP_lucwS6w"
      },
      "outputs": [],
      "source": [
        "class BetterModel:\n",
        "\n",
        "  def __init__(self):\n",
        "    self.bias = tf.Variable(0.)\n",
        "    self.weight = tf.Variable(2.)\n",
        "\n",
        "@tf.function\n",
        "def evaluate(model, x):\n",
        "  return model.weight * x + model.bias\n",
        "\n",
        "better_model = BetterModel()\n",
        "print(evaluate(better_model, x))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ktqwMJBqwTFj"
      },
      "outputs": [],
      "source": [
        "print(\"Adding bias!\")\n",
        "better_model.bias.assign_add(5.0)  # Note: instead of better_model.bias += 5\n",
        "print(evaluate(better_model, x))  # This works!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lPr_6mK_AQWL"
      },
      "source": [
        "### tf.Variables の作成\n",
        "\n",
        "`Function` では、最初に呼び出された時に 1 回だけ変数を作成し、以降ではそれが再利用されます。新しいトレースで `tf.Variables` を作成することはできません。以降の呼び出して新しい変数を作成することはできませんが、将来的には可能になる予定です。\n",
        "\n",
        "例:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Tx0Vvnb_9OB-"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def f(x):\n",
        "  v = tf.Variable(1.0)\n",
        "  return v\n",
        "\n",
        "with assert_raises(ValueError):\n",
        "  f(1.0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KYm6-5GCILXQ"
      },
      "source": [
        "関数が初めて実行される際に作成される変数である限り、それらの変数を `Function` 内で作成できます。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HQrG5_kOiKl_"
      },
      "outputs": [],
      "source": [
        "class Count(tf.Module):\n",
        "  def __init__(self):\n",
        "    self.count = None\n",
        "\n",
        "  @tf.function\n",
        "  def __call__(self):\n",
        "    if self.count is None:\n",
        "      self.count = tf.Variable(0)\n",
        "    return self.count.assign_add(1)\n",
        "\n",
        "c = Count()\n",
        "print(c())\n",
        "print(c())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7uD6qI7aJwbR"
      },
      "source": [
        "#### 複数の Keras オプティマイザとの使用\n",
        "\n",
        "2 つ以上の Keras オプティマイザを `tf.function` で使用しようとすると、「`ValueError: tf.function-decorated function tried to create variables on non-first call.`」というエラーが発生することがあります。このエラーは、オプティマイザが初めて勾配を適用する際に、内部的に `tf.Variables` を作成するために発生するものです。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yWQ3-r99Jvze"
      },
      "outputs": [],
      "source": [
        "opt1 = tf.keras.optimizers.Adam(learning_rate = 1e-2)\n",
        "opt2 = tf.keras.optimizers.Adam(learning_rate = 1e-3)\n",
        " \n",
        "@tf.function\n",
        "def train_step(w, x, y, optimizer):\n",
        "   with tf.GradientTape() as tape:\n",
        "       L = tf.reduce_sum(tf.square(w*x - y))\n",
        "   gradients = tape.gradient(L, [w])\n",
        "   optimizer.apply_gradients(zip(gradients, [w]))\n",
        "\n",
        "w = tf.Variable(2.)\n",
        "x = tf.constant([-1.])\n",
        "y = tf.constant([2.])\n",
        "\n",
        "train_step(w, x, y, opt1)\n",
        "print(\"Calling `train_step` with different optimizer...\")\n",
        "with assert_raises(ValueError):\n",
        "  train_step(w, x, y, opt2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7Q8BRPCThTjB"
      },
      "source": [
        "トレーニング中にオプティマイザを変更する必要がある場合は、回避策として、オプティマイザごとに新しい `Function` を作成し、[`ConcreteFunction`](#obtaining_concrete_functions) を直接呼び出すようにすることができます。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YV5F2Gy9hSI3"
      },
      "outputs": [],
      "source": [
        "opt1 = tf.keras.optimizers.Adam(learning_rate = 1e-2)\n",
        "opt2 = tf.keras.optimizers.Adam(learning_rate = 1e-3)\n",
        "\n",
        "# Not a tf.function.\n",
        "def train_step(w, x, y, optimizer):\n",
        "   with tf.GradientTape() as tape:\n",
        "       L = tf.reduce_sum(tf.square(w*x - y))\n",
        "   gradients = tape.gradient(L, [w])\n",
        "   optimizer.apply_gradients(zip(gradients, [w]))\n",
        "\n",
        "w = tf.Variable(2.)\n",
        "x = tf.constant([-1.])\n",
        "y = tf.constant([2.])\n",
        "\n",
        "# Make a new Function and ConcreteFunction for each optimizer.\n",
        "train_step_1 = tf.function(train_step).get_concrete_function(w, x, y, opt1)\n",
        "train_step_2 = tf.function(train_step).get_concrete_function(w, x, y, opt2)\n",
        "for i in range(10):\n",
        "  if i % 2 == 0:\n",
        "    train_step_1(w, x, y) # `opt1` is not used as a parameter. \n",
        "  else:\n",
        "    train_step_2(w, x, y) # `opt2` is not used as a parameter."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Xjnz5CcuqQac"
      },
      "source": [
        "#### 複数の Keras モデルとの使用\n",
        "\n",
        "また、別のモデルインスタンスを同一の  `Function` に渡す際に、「`ValueError: tf.function-decorated function tried to create variables on non-first call.`」というエラーも発生することがあります。\n",
        "\n",
        "このエラーは、Keras モデル（[入力形状が定義されていない](https://www.tensorflow.org/guide/keras/custom_layers_and_models#best_practice_deferring_weight_creation_until_the_shape_of_the_inputs_is_known)）と Keras レイヤーが、初めて呼び出されるときに `tf.Variables` を作成するために発生するものです。これらの変数をすでに呼び出された `Function` 内で初期化しようとしているのでしょう。このエラーを回避するには、モデルをトレーニングする前に、`model.build(input_shape)` を呼び出して、すべての重みを初期化するようにしてください。\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IKyrEY5GVX3M"
      },
      "source": [
        "## 参考資料\n",
        "\n",
        "`Function` のエクスポートと読み込みの方法については、[SavedModel ガイド](../../guide/saved_model)をご覧ください。トレーシングの後に実行するグラフの最適化については、[Grappler ガイド](../../guide/graph_optimization)をご覧ください。データパイプラインの最適化方法とモデルのプロファイリングについては、[Profiler ガイド](../../guide/profiler.md)をご覧ください。"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "function.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
